# 构建一个简单的计算器Agent

from dotenv import load_dotenv
import os
# from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from langchain_core.tools import tool # 用于@tool装饰器
from langchain.agents import AgentExecutor, create_react_agent # 核心的Agent执行器
from langchain.agents.format_scratchpad import format_to_openai_function_messages # 格式化中间步骤
from langchain.agents.output_parsers import OpenAIFunctionsAgentOutputParser # 解析LLM输出
from langchain_core.runnables import RunnablePassthrough
from langchain_core.messages import AIMessage
from langchain.agents.format_scratchpad import format_log_to_str

# https://python.langchain.com/docs/how_to/agent_executor/
# recommend use langgraph create_react_agent instead 

def format_ollama_messages(steps):
    return [
        AIMessage(content=str(action) + "\n" + str(observation))
        for action, observation in steps
    ]

load_dotenv()
# llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)

from langchain_ollama import ChatOllama
llm = ChatOllama(model="qwen3:8b", temperature=0.5, reasoning=False)

# 定义工具 (复用上面定义的 add 和 multiply)
@tool
def add(a: int, b: int) -> int:    
    """Adds two integers and returns the result."""    
    return a + b

@tool
def multiply(a: int, b: int) -> int:    
    """Multiplies two integers and returns the result."""    
    return a * b

tools = [add, multiply] # 将工具列表传入

# 1. 定义 Agent 的提示模板
# # MessagesPlaceholder("agent_scratchpad") 是关键，它会插入LLM的思考过程和工具执行结果
prompt = ChatPromptTemplate.from_messages([    
    ("system", "你是一个数学助手，可以使用工具进行加法和乘法运算。"),    
    MessagesPlaceholder(variable_name="chat_history"), # 可以选择性加入聊天历史    
    ("user", "{input}"),    
    MessagesPlaceholder(variable_name="agent_scratchpad"), # LLM的思考过程和工具输出会填充这里
])

# 2. 将LLM绑定到工具 (这是 OpenAI Function Calling 的核心)
# LLM 会自动知道如何使用这些工具，并在需要时生成工具调用请求
# llm_with_tools = llm.bind_tools(tools)
# llm_with_tools = llm.bind(functions=tools)

# 3. 构建 Agent 核心逻辑 (LCEL 链)
# 这一步是关键！它定义了 Agent 的思考循环
# agent_runnable = RunnablePassthrough.assign(    
#     agent_scratchpad=lambda x: format_to_openai_function_messages(x["intermediate_steps"])
# ) | prompt | llm_with_tools | OpenAIFunctionsAgentOutputParser()

# agent_runnable = {
#     "input": lambda x: x["input"],
#     "chat_history": lambda x: x["chat_history"], # 如果需要聊天历史，可以在这里传入，否则可以省略这一行和 chat_history 参数
#     "agent_scratchpad": lambda x: format_log_to_str(x["intermediate_steps"])
#     # "agent_scratchpad": lambda x: [
#     #     f"工具调用: {action.tool}\n输入: {action.tool_input}\n结果: {observation}"
#     #     for action, observation in x["intermediate_steps"]
#     # ]
# } | prompt | llm_with_tools

agent_runnable = create_react_agent(llm, tools, prompt)

# 4. 创建 AgentExecutor
# AgentExecutor 是实际运行 Agent 的组件，它会处理 LLM 的响应 (是文本还是工具调用)，
# 如果是工具调用，它会执行工具并把结果反馈给 LLM，直到任务完成。
agent_executor = AgentExecutor(agent=agent_runnable, tools=tools, verbose=True) # verbose=True 会打印详细的思考过程
print("--- 简单的计算器 Agent 示例 ---")
print("Agent 正在运行中...")

# 调用 Agent
response_calc1 = agent_executor.invoke({"input": "256乘以48是多少？", "chat_history": []})
print(f"问题: 256乘以48是多少？")
print(f"Agent 回答: {response_calc1['output']}\n")

# response_calc2 = agent_executor.invoke({"input": "200加300再乘以2是多少？", "chat_history": []})
# print(f"问题: 200加300再乘以2是多少？")
# print(f"Agent 回答: {response_calc2['output']}\n")

# response_no_tool = agent_executor.invoke({"input": "你好，能帮我做点什么？", "chat_history": []})
# print(f"问题: 你好，能帮我做点什么？")
# print(f"Agent 回答: {response_no_tool['output']}\n") # 不需要工具，直接回答

